# -*- coding:utf-8 -*-
import os
import os.path as osp
import fnmatch
from tool import filesystem
import shutil
from tqdm import tqdm
import os, glob
from scipy.io import loadmat
from collections import defaultdict
import numpy as np
from lxml import etree, objectify

def open_save(file, savepath):
    f = open(file, 'rb')
    # 将seq文件的内容转化成str类型
    string = f.read().decode('latin-1')

    # splitstring是图片的前缀，可以理解成seq是以splitstring为分隔的多个jpg合成的文件
    splitstring = "\xFF\xD8\xFF\xE0\x00\x10\x4A\x46\x49\x46"  
    # split函数做一个测试,因此返回结果的第一个是在seq文件中是空，因此后面省略掉第一个
    strlist = string.split(splitstring)

    f.close()
    count = 0  

    # 遍历每一个jpg文件内容，然后加上前缀合成图片
    for img in strlist:        
        filename = str(count)+'.jpg'
        filename = osp.join(savepath , filename)
        if count > 0:                                    
            i = open(filename, 'wb+')
            i.write(splitstring.encode('latin-1'))
            i.write(img.encode('latin-1'))
            i.close()
        count = count + 1
        print("Generating JPEGImages jpg file of picture:{}".format(filename))


def vbb_anno2dict(vbb_file, cam_id):
    filename = os.path.splitext(os.path.basename(vbb_file))[0]

    # 定义字典对象annos
    annos = defaultdict(dict)
    vbb = loadmat(vbb_file)
    # object info in each frame: id, pos, occlusion, lock, posv
    objLists = vbb['A'][0][0][1][0]
    objLbl = [str(v[0]) for v in vbb['A'][0][0][4][0]]     # 可查看所有类别
    # person index
    person_index_list = np.where(np.array(objLbl) == "person")[0]   # 只选取类别为‘person’的xml
    for frame_id, obj in enumerate(objLists):
        if len(obj) > 0:
            frame_name = str(cam_id) + "_" + str(filename) + "_" + str(frame_id+1) + ".jpg"
            annos[frame_name] = defaultdict(list)
            annos[frame_name]["id"] = frame_name
            annos[frame_name]["label"] = "person"
            for id, pos, occl in zip(obj['id'][0], obj['pos'][0], obj['occl'][0]):
                id = int(id[0][0]) - 1  # for matlab start from 1 not 0
                if not id in person_index_list:  # only use bbox whose label is person
                    continue
                pos = pos[0].tolist()
                occl = int(occl[0][0])
                annos[frame_name]["occlusion"].append(occl)
                annos[frame_name]["bbox"].append(pos)
            if not annos[frame_name]["bbox"]:
                del annos[frame_name]
    print(annos)
    return annos


def instance2xml_base(anno, bbox_type='xyxy'):
    """bbox_type: xyxy (xmin, ymin, xmax, ymax); xywh (xmin, ymin, width, height)"""
    assert bbox_type in ['xyxy', 'xywh']
    E = objectify.ElementMaker(annotate=False)
    anno_tree = E.annotation(
        E.folder('VOC2014_instance/person'),
        E.filename(anno['id']),
        E.source(
            E.database('Caltech pedestrian'),
            E.annotation('Caltech pedestrian'),
            E.image('Caltech pedestrian'),
            E.url('None')
        ),
        E.size(
            E.width(640),
            E.height(480),
            E.depth(3)
        ),
        E.segmented(0),
    )
    for index, bbox in enumerate(anno['bbox']):
        bbox = [float(x) for x in bbox]
        if bbox_type == 'xyxy':
            xmin, ymin, w, h = bbox
            xmax = xmin+w
            ymax = ymin+h
        else:
            xmin, ymin, xmax, ymax = bbox
        E = objectify.ElementMaker(annotate=False)
        anno_tree.append(
            E.object(
            E.name(anno['label']),
            E.bndbox(
                E.xmin(int(xmin)),
                E.ymin(int(ymin)),
                E.xmax(int(xmax)),
                E.ymax(int(ymax))
            ),
            E.difficult(0),
            E.occlusion(anno["occlusion"][index])
            )
        )
    return anno_tree


def parse_anno_file(vbb_inputdir, vbb_outputdir):
    # annotation sub-directories in hda annotation input directory
    assert os.path.exists(vbb_inputdir)
    sub_dirs = os.listdir(vbb_inputdir)     # 对应set00,set01...
    for sub_dir in sub_dirs:
        print("Parsing annotations of camera: ", sub_dir)
        cam_id = sub_dir
        # 获取某一个子set下面的所有vbb文件
        vbb_files = glob.glob(os.path.join(vbb_inputdir, sub_dir, "*.vbb"))
        for vbb_file in vbb_files:
            # 返回一个vbb文件中所有的帧的标注结果
            annos = vbb_anno2dict(vbb_file, cam_id)
            
            if annos:
                # 组成xml文件的存储文件夹，形如“/Users/chenguanghao/Desktop/Caltech/xmlresult/”
                vbb_outdir = vbb_outputdir
                                              
                # 如果不存在 vbb_outdir
                if not os.path.exists(vbb_outdir):
                    os.makedirs(vbb_outdir)

                for filename, anno in sorted(annos.items(), key=lambda x: x[0]):                  
                    if "bbox" in anno:
                        anno_tree = instance2xml_base(anno)
                        outfile = os.path.join(vbb_outdir, os.path.splitext(filename)[0]+".xml")
                        print("Generating annotation xml file of picture: ", filename)
                        # 生成最终的xml文件，对应一张图片
                        etree.ElementTree(anno_tree).write(outfile, pretty_print=True)


def visualize_bbox(xml_file, img_file):
    import cv2
    tree = etree.parse(xml_file)
    # load image
    image = cv2.imread(img_file)
    origin = cv2.imread(img_file)
    # 获取一张图片的所有bbox
    for bbox in tree.xpath('//bndbox'):
        coord = []
        for corner in bbox.getchildren():
            coord.append(int(float(corner.text)))
        print(coord)
        cv2.rectangle(image, (coord[0], coord[1]), (coord[2], coord[3]), (0, 0, 255), 2)
    # visualize image
    cv2.imshow("test", image)
    cv2.imshow('origin', origin)
    cv2.waitKey(0)




def last():

    xml_path = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/Annotations_"
    jpg_path = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/images"
    save_path = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/images_"


    filter = []
    print("filter loading...")
    for file in tqdm(os.listdir(xml_path)):
        filter.append(os.path.splitext(file)[0])

    for file in os.listdir(jpg_path):
        if os.path.splitext(file)[0] in filter:
            old_path = os.path.join(jpg_path, file)
            new_path = os.path.join(save_path, file)
            shutil.copyfile(old_path, new_path)
            print("Success save jpg file: {}".format(new_path))


def filter_xml():
    space = 40*70
    xml_path = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/Annotations"
    save_path = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/Annotations_"

    for file_name in os.listdir(xml_path):
        file_path = os.path.join(xml_path, file_name)
        print(file_path)
        in_file = open(file_path)
        tree = ET.parse(in_file)  # ET是一个xml文件解析库，ET.parse（）打开xml文件。parse--"解析"
        root = tree.getroot()  # 获取根节点

        write_flag = False

        for obj in root.findall('object'):  # 找到根节点下所有“object”节点
            bbox = obj.find('bndbox')
            x1 = int(bbox.find('xmin').text)
            y1 = int(bbox.find('ymin').text)
            x2 = int(bbox.find('xmax').text)
            y2 = int(bbox.find('ymax').text)

            if (x2-x1)*(y2-y1) >= space:
                write_flag = True
        if write_flag:
            tree.write(os.path.join(save_path, file_name))
            print("save file: {}".format(os.path.join(save_path, file_name)))



def seq2img():
    seq_inputdir = "/home/xc/Downloads/set"

    for filepath in filesystem.get_all_filepath(seq_inputdir):
        # 若果是set01、set02……这类文件夹
        save_dir = ".".join(filepath.split(".")[:-1])
        os.makedirs(save_dir, exist_ok=True)
        print(save_dir)
        open_save(filepath, save_dir)


def vbb2xml():
    vbb_inputdir = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/data/annotations"
    vbb_outputdir = "/home/zhou/数据集/行人检测/Caltech_Pedestrian_Detection_Benchmark/Annotations"
    parse_anno_file(vbb_inputdir, vbb_outputdir)


if __name__ == '__main__':
    seq2img()

    # vbb2xml()

    # filter_xml()

    # last()